import redpic as rp
import kenv as kv
import holoviews as hv
import numpy as np
import pandas as pd
import glob
hv.notebook_extension('matplotlib')
%output size=150 backend='matplotlib' fig='png' dpi=200
%opts Curve [show_grid=True aspect=5] (linewidth=1 alpha=0.7 color='blue')
%opts Scatter [show_grid=True aspect=5] (alpha=0.7 s=0.5)
kenbeam = kv.Beam(energy=2, current=0, radius=55e-3, rp=.0e0, normalized_emittance=57e-6)
redbeam = rp.Beam(rp.electron) # off space charge
KV = rp.Distribution(name='KV', x=0.055, y=0.055, z=2.5, px=0.0, py=0.0, pz=2.4585) # KV distribution
redbeam.generate(KV, n=1e4)
acc = kv.Accelerator(0.0, 6.7, 0.01)
Ez_beamline = {}
for z0, E0, filename, name in [
# m MV/m Unique name
[ 4.223, -1.2, 'Ez.dat', 'Acc. 1'],
[ 6.076, -1.2, 'Ez.dat', 'Acc. 2'],
]:
acc.Ez_beamline[name] = kv.Element(z0, E0, filename, name)
Bz_beamline = {}
for z0, B0, filename, name in [
# m T Unique name
[ 0.950, 0.035, 'Bz.dat', 'Sol. 1'],
[ 2.100, 0.032, 'Bz.dat', 'Sol. 2'],
[ 2.900, 0.043, 'Bz.dat', 'Sol. 3'],
[ 3.800, 0.038, 'Bz.dat', 'Sol. 4'],
[ 4.700, 0.047, 'Bz.dat', 'Sol. 5'],
[ 5.600, 0.052, 'Bz.dat', 'Sol. 6'],
[ 6.500, 0.050, 'Bz.dat', 'Sol. 7'],
]:
acc.Bz_beamline[name] = kv.Element(z0, B0, filename, name)
acc.compile()
kensim = kv.Simulation(kenbeam, acc)
redsim = rp.Simulation(redbeam, acc)
kensim.track()
redsim.track()
track_files = np.sort(glob.glob('*.*[0-9].csv'))
cols = ['x', 'y', 'z', 'px', 'py', 'pz', 'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']
i_to_plot = np.arange(0, len(track_files), 1)
def read_track(fname):
df = pd.read_csv(fname, dtype='float32')
df['x'] = df['x']*1e3 # mm
df['y'] = df['y']*1e3 # mm
df['Bx'] = df['Bx']*1e4 # Gs
df['By'] = df['By']*1e4 # Gs
df['Bz'] = df['Bz']*1e4 # Gs
#df = df.sample(n=100000)
return df
dim_x = hv.Dimension('x', unit='mm', range=(-100, 100))
dim_y = hv.Dimension('y', unit='mm')
dim_z = hv.Dimension('z', unit='m', range=(0, 5))
dim_px = hv.Dimension('px', unit='MeV/c', label='$p_x$')
dim_py = hv.Dimension('py', unit='MeV/c', label='$p_y$')
dim_pz = hv.Dimension('pz', unit='MeV/c', label='$p_z$')
dim_Ex = hv.Dimension('Ex', unit='MV/m', label='$E_x$')
dim_Ey = hv.Dimension('Ey', unit='MV/m', label='$E_y$')
dim_Ez = hv.Dimension('Ez', unit='MV/m', label='$E_z$')
dim_Bx = hv.Dimension('Bx', unit='Gs', label='$B_x$')
dim_By = hv.Dimension('By', unit='Gs', label='$B_y$')
dim_Bz = hv.Dimension('Bz', unit='Gs', label='$B_z$')
def plot(i):
fname = track_files[i]
df = read_track(fname)
ken_z_x = hv.Curve(((acc.z,kensim.envelope_x(acc.z)*1e3)), kdims=[dim_z], vdims=[dim_x], label='kenv')*\
hv.Curve(((acc.parameter,-kensim.envelope_x(acc.z)*1e3)), kdims=[dim_z], vdims=[dim_x])
red_z_x = hv.Scatter(df, kdims=[dim_z, dim_x], label='redpic')
red_z_px = hv.Scatter(df, kdims=[dim_z, dim_px], label='redpic')
red_z_pz = hv.Scatter(df, kdims=[dim_z, dim_pz], label='redpic')
red_z_Ez = hv.Scatter(df, kdims=[dim_z, dim_Ez], label='redpic')
red_z_Bz = hv.Scatter(df, kdims=[dim_z, dim_Bz], label='redpic')
ken_z_Bz = hv.Curve(((acc.z, acc.Bz(acc.z)*1e4)), kdims=[dim_z], vdims=[dim_Bx], label='kenv')
return (red_z_x*ken_z_x + red_z_Bz*ken_z_Bz).cols(1)
items = [(i, plot(i)) for i in i_to_plot]
hv.HoloMap(items, kdims = ['Track file index']).collate()